import csv
import sys
import os
import kornia
import matplotlib
import matplotlib.pyplot as plt
import numpy
import torch
from PyQt5.QtGui import QMovie
from torch import Tensor

matplotlib.use("Qt5Agg")
plt.rcParams['font.family'] = ['sans-serif']
plt.rcParams['font.sans-serif'] = ['SimHei']

import numpy as np
from PyQt5.QtWidgets import QWidget, QApplication, QMainWindow, QVBoxLayout, QLabel, QInputDialog
# from cicle import Ui_MainWindow
from PyQt5.QtWidgets import QApplication, QSizePolicy, QFileDialog, QTableWidgetItem
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5 import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
from matplotlib.patches import Ellipse, Circle, Wedge
from PyQt5 import QtCore, QtGui, QtWidgets

from demo import Ui_MainWindow
from PIL import Image
from Main_Decoder import Decoder
from Main_Encoder import Encoder
from util.Huffman_Tree import Huffman_Tree
from util.Rscode_c import Rscode_c
from PyQt5.Qt import pyqtSignal


class MyFigure(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=500):
        # https://blog.csdn.net/The_Time_Runner/article/details/89312660 画图的链接
        # https://blog.csdn.net/weixin_43988887/article/details/123590179
        # plt.rcParams['font.family'] = ['SimHei']
        # plt.rcParams['axes.unicode_minus'] = False
        # 创建一个2*2布局的表格
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        # 设置横纵坐标的名称以及对应字体格式
        FigureCanvas.__init__(self, self.fig)
        self.setParent(parent)

        # FigureCanvas.setSizePolicy(self, QSizePolicy.Expanding, QSizePolicy.Expanding)
        # FigureCanvas.updateGeometry(self)
        self.fig.tight_layout()

    def draw_single(self, path, title=""):
        """
        画单张图
        :param point: 二维数组,表示点的坐标 [[0.43, -0.9], [1.94, 0.5]]
        :param idx:
        :return:
        """
        self.fig.clear()
        ax = self.fig.add_subplot(111)
        image = Image.open(path).convert("RGB")
        image = np.array(image)
        H, W, C = image.shape
        ax.set_title("{} H={} W={}".format(title, H, W))
        ax.imshow(image)
        self.draw()

    def draw_tensor(self, image: Tensor, title=""):
        """
        画单张图
        :param point: 二维数组,表示点的坐标 [[0.43, -0.9], [1.94, 0.5]]
        :param idx:
        :return:
        """
        self.fig.clear()
        ax = self.fig.add_subplot(111)
        image = image.permute(0, 2, 3, 1).numpy().squeeze()
        print("image", image.shape)
        H, W, C = image.shape
        ax.set_title("{} H={} W={}".format(title, H, W))
        ax.imshow(image)
        self.draw()


class M_mask(QWidget):

    def mouseMoveEvent(self, QMouseEvent):
        QMouseEvent.accept()

    sig_resize = pyqtSignal(int, int)

    def __init__(self, l="", h=""):
        super(M_mask, self).__init__()
        if l == "":
            l = 640
        if h == "":
            h = 512
        self.resize(l, h)
        self.m_ui()
        self.raise_()
        # QtCore.Qt.WindowFrameSection

        self.setWindowFlag(QtCore.Qt.FramelessWindowHint)  # 隐藏边框
        self.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)  # 置顶

    def m_ui(self):
        self.setWindowTitle("加载gif动画")
        self.label_01()

    def label_01(self):
        self.label = QLabel(self)
        self.label.setScaledContents(True)
        self.m_movie()

    def m_movie(self):
        self.movie = QMovie("loading.gif")
        self.label.setMovie(self.movie)
        self.movie.start()

    def resizeEvent(self, event):
        self.label.resize(self.width(), self.height())


class BilibiliKit(QMainWindow, Ui_MainWindow):
    def __init__(self, parent=None):
        super(BilibiliKit, self).__init__(parent)

        #  添加绘图主键
        self.setupUi(self)

        #  添加绘图主键
        self.origin = MyFigure(width=16, height=12, dpi=80)
        origin_toolbar = NavigationToolbar(self.origin, self)
        self.origin_gridLayout.addWidget(origin_toolbar)
        self.origin_gridLayout.addWidget(self.origin)

        self.hidden = MyFigure(width=16, height=12, dpi=80)
        hidden_toolbar = NavigationToolbar(self.hidden, self)
        self.hidden_gridLayout.addWidget(hidden_toolbar)
        self.hidden_gridLayout.addWidget(self.hidden)

        # self.loading = M_mask()
        # self.loading.show()

        #  添加信息
        self.myretranslateUi(parent)
        self.connectSlots()

        self.image_path = None
        self.hidden_path = None
        self.hidden_image_path = None
        self.save_path = None
        self.ratio = 1

        # 2 表示正在加载csv

    def open_image_file(self):
        """打开一个图像文件"""
        # filter="jpg Files (*.jpg);;png Files (*.png);;All Files (*)"# 弹出窗口的后缀名过滤器
        fileName, filetype = QFileDialog.getOpenFileName(self,
                                                         caption="open file",  # 弹出窗口的窗口名
                                                         directory='.',  # 弹出窗口的默认目录，不填默认为项目的根目录
                                                         filter="png,jpg (*.png *.jpg *.jpeg);;All Files (*)"
                                                         # 弹出窗口的后缀名过滤器
                                                         )
        return fileName

    def open_hidden_file(self):
        """打开一个图像文件"""
        # filter="jpg Files (*.jpg);;png Files (*.png);;All Files (*)"# 弹出窗口的后缀名过滤器
        fileName, filetype = QFileDialog.getOpenFileName(self,
                                                         caption="open file",  # 弹出窗口的窗口名
                                                         directory='.',  # 弹出窗口的默认目录，不填默认为项目的根目录
                                                         filter="txt,py (*.txt *.py);;All Files (*)"
                                                         # 弹出窗口的后缀名过滤器
                                                         )
        return fileName

    def myretranslateUi(self, MainWindow):
        """
        初始化信息
        :param MainWindow:
        :return:
        """
        _translate = QtCore.QCoreApplication.translate

    def connectSlots(self):
        """
        点击时间集合
        :return:
        """
        # self.index.currentIndexChanged[int].connect(self.index_change)  # 条目发生改变，发射信号，传递条目索引
        # self.rule.currentIndexChanged[int].connect(self.rule_change)  # 条目发生改变，发射信号，传递条目索引
        #
        # self.open.clicked.connect(self.open_csv)  # 点击打开时间
        # self.reset.clicked.connect(self.reset_click)
        self.attack.clicked.connect(self.attack_click)
        self.open_image.clicked.connect(self.open_image_click)
        self.open_hidden.clicked.connect(self.open_hidden_click)
        self.encryption.clicked.connect(self.encryption_click)
        self.open_hidden_image.clicked.connect(self.open_hidden_image_click)

        self.decrypt.clicked.connect(self.decrypt_click)

    def decrypt_click(self):
        """
        开始解密
        """
        try:
            if os.path.exists(self.hidden_image_path):
                # self.hidden.draw_single(self.hidden_image_path, title="隐写图像")
                d = Decoder(self.hidden_image_path, 1)
                decoded_messages_raw = d.network.encoder_decoder.decoder(d.img_decode_tensor)
                self.hidden.draw_tensor(d.img_decode_tensor, "被攻击后图像大小")
                decoded_messages = d.round_0_1_(decoded_messages_raw)[0].tolist()
                file_name, code_len, res_len = d.get_txt_code_file_and_len(decoded_messages)
                # rs 解码消息
                rs_c = Rscode_c()
                original_message = rs_c.decode(decoded_messages[0:code_len])
                original_message = original_message[0:len(original_message) - res_len]
                tree = Huffman_Tree(None)
                m = tree.decode_c(os.path.join("file", str(file_name) + ".txt"), original_message,
                                  os.path.join("result", str(file_name) + ".py"))
                self.hidden_textBrowser.setText(m)
                filepath = self.save_txt()
                with open(filepath, "a", encoding="utf-8") as f:
                    f.write(m)
            else:
                self.show_label("隐写图片路径不存在")

        except Exception as e:
            self.show_label("{}".format(e))

    def encryption_click(self):
        """
        开始加密
        """
        image_path = self.image_path
        hidden_path = self.hidden_path
        try:
            if os.path.exists(image_path) and os.path.exists(hidden_path):
                e = Encoder(image_path, hidden_path)
                with torch.no_grad():
                    # use device to compute
                    images, messages = e.cover_tensor_ready.to(e.network.device), e.message_tensor_ready.to(
                        e.network.device)
                    encoded_images = e.network.encoder_decoder.encoder(images, messages)
                    # print("encoded_images", encoded_images.shape)
                    ssim = kornia.losses.ssim_loss(encoded_images.detach(), images, 5).item()
                    # PSNR的全称为“Peak Signal-to-Noise Ratio”，直译为中文就是峰值信噪比。是一种衡量图像质量的指标。
                    # 在很多领域都会需要这个指标，比如在超分辨率重建图像的时候，PSNR就是很重要的指标了。
                    # 峰值信噪比（英语：Peak signal-to-noise ratio，常缩写为PSNR）是一个表示信号最大可能功率和影响它的表示精度的破坏性噪声功率的比值的工程术语。
                    # 由于许多信号都有非常宽的动态范围，峰值信噪比常用对数分贝单位来表示。
                    # PSNR高于40dB说明图像质量极好（即非常接近原始图像），
                    # 在30—40dB通常表示图像质量是好的（即失真可以察觉但可以接受），
                    # 在20—30dB说明图像质量差；
                    # 最后，PSNR低于20dB图像不可接受
                    psnr = kornia.losses.psnr_loss(encoded_images.detach(), images, 2).item()
                    self.show_label("评价指标 psnr={:.3f},SSIM={:.3f}".format(ssim, psnr))
                    filepath = self.save_png()
                    e.save_encode_image(filepath, encoded_images)
                    #  设置解密信息
                    self.hidden_image_path = filepath
            else:
                if self.image_path is None or self.hidden_path is None:
                    self.show_label("请添加文件，再进行加密")
                else:
                    self.show_label("文件路径不存在")

        except Exception as e:
            self.show_label("{}".format(e))

    def save_txt(self):
        filepath, type = QFileDialog.getSaveFileName(self, "文件保存", "/",
                                                     'txt(*.txt *.py)')  # 前面是地址，后面是文件类型,得到输入地址的文件名和地址txt(*.txt*.xls);;image(*.png)不同类别

        print(filepath)
        return filepath

    def save_png(self):
        filepath, type = QFileDialog.getSaveFileName(self, "文件保存", "/",
                                                     'jpg(*.jpg *.png *.jpeg)')  # 前面是地址，后面是文件类型,得到输入地址的文件名和地址txt(*.txt*.xls);;image(*.png)不同类别
        # self.show_label()
        print(filepath)
        return filepath

    def showDialog(self):
        text, ok = QInputDialog.getText(self, 'Input Dialog','请输入缩放比例0.9-1.1之间:')
        if ok:
            return float(text)
        else:
            return 1

    def attack_click(self):
        try:
            ratio = self.showDialog()
            if isinstance(ratio, float) or isinstance(ratio, int):
                self.ratio = ratio
                self.show_label("攻击成功，图片缩放*{}".format(ratio))
        except Exception as e:
            self.show_label("{}".format(e))

    # 点击了攻击的按钮
    # pass

    def open_hidden_image_click(self):
        try:
            fileName = self.open_image_file()
            self.show_label(fileName)
            self.hidden_image_path = fileName
            self.hidden.draw_single(fileName, "可能含有信息的隐写图")
        except Exception as e:
            self.show_label("{}".format(e))

    def show_label(self, text):
        self.information.setText(text)

    def open_hidden_click(self):
        """
        打开隐写文件
        """
        try:
            hidden_path = self.open_hidden_file()
            self.hidden_path = hidden_path
            with open(hidden_path, "r", encoding="utf-8") as f:
                self.origin_textBrowser.setText(f.read())
        except Exception as e:
            self.show_label("{}".format(e))

    def open_image_click(self):
        """
        打开原图
        """
        try:
            fileName = self.open_image_file()
            self.show_label(fileName)
            self.image_path = fileName
            self.origin.draw_single(fileName, "origin image ")
        except Exception as e:
            self.show_label("{}".format(e))


if __name__ == '__main__':
    app = QApplication(sys.argv)
    w = BilibiliKit()
    w.show()
    sys.exit(app.exec_())
